import typing
from fastapi import APIRouter
from starlette.endpoints import WebSocketEndpoint
from starlette.websockets import WebSocket,WebSocketDisconnect

class DashboardWS(WebSocketEndpoint):
    # json None
    encoding = None
    active_connections = []
    async def on_connect(self, ws:WebSocket) -> None:
        await ws.accept()  # 完成握手
        print("连接成功")
        # u_type = web_socket.query_params.get("u_type")
        # token = web_socket.headers.get("sec-websocket-protocol")
        # real_ip = web_socket.headers.get('x-forwarded-for')
        # real_host = web_socket.headers.get("host")
        try:
            self.active_connections.append(ws)
        except ValueError:
            pass

    async def on_receive(self, ws: WebSocket, data: typing.Any) -> None:
        try:
            if isinstance(data, str):
                # 处理文本消息
                print(f"Received text message: {data}")
                await ws.send_text(f"Echo: {data}")
            elif isinstance(data, bytes):
                # 处理二进制消息
                print(f"Received binary message: {data}")
                await ws.send_bytes(data)
            else:
                print(f"Unknown message type: {data}")
        except Exception as e:
            print(e)

    async def on_disconnect(self, ws: WebSocket, close_code: int) -> None:
        print("断开连接")
        # 执行逻辑，从对象池，删除用户
        self.active_connections.remove(ws)
        pass

router = APIRouter()
router.add_websocket_route('/ws', DashboardWS)

